package com.lrj.datastream.transform.partition;

import com.lrj.datastream.domain.ClickEvent;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

/**
 * @author lrj
 * @date 2022/3/29 10:06
 */
public class TestPartition {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<ClickEvent> stream = env.fromElements(
                new ClickEvent("Mary", "./cart", 1000L),
                new ClickEvent("Bob", "./home", 1000L),
                new ClickEvent("Mary", "./home", 1000L),
                new ClickEvent("Alice", "./product?id=10", 1000L)
        );
        // 1.随机分区,默认,没有指定分区器
        stream.shuffle().print().setParallelism(4);
        // 2.轮询分区,重平衡
        stream.rebalance().print("re balance").setParallelism(4);
        // 3.重缩放分区
        DataStreamSource<Integer> intStream = env.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> ctx) throws Exception {
                for (int i = 0; i < 8; i++) {
                    if (getRuntimeContext().getIndexOfThisSubtask() == i % 2) {
                        ctx.collect(i);
                    }
                }
            }

            @Override
            public void cancel() {

            }
        }).setParallelism(2);
        intStream.print("before").setParallelism(2);
        // 分组内重分区
        intStream.rescale().print("rescale").setParallelism(4);
        // 所有进行重分区
        intStream.rebalance().print("rebalance").setParallelism(4);
        /*
        broadcast:1> 1
        broadcast:4> 1
        broadcast:3> 1
        broadcast:3> 3
        broadcast:2> 1
        broadcast:3> 5
        broadcast:4> 3
        broadcast:1> 3
        broadcast:4> 5
        broadcast:3> 7
        broadcast:2> 3
        broadcast:4> 7
        broadcast:1> 5
        broadcast:2> 5
        broadcast:1> 7
        broadcast:2> 7
        broadcast:1> 0
        broadcast:1> 2
        broadcast:3> 0
        broadcast:3> 2
        broadcast:1> 4
        broadcast:4> 0
        broadcast:1> 6
        broadcast:3> 4
        broadcast:3> 6
        broadcast:4> 2
        broadcast:4> 4
        broadcast:4> 6
        broadcast:2> 0
        broadcast:2> 2
        broadcast:2> 4
        broadcast:2> 6
         */
        // 广播到所有分区去,会导致所有分区数据都一样
        intStream.broadcast().print("broadcast").setParallelism(4);
        // 全局分区,将所有数据分配到一个分区上去
        intStream.global().print("global").setParallelism(4);
        /*
        custom:1> 3
        custom:1> 6
        custom:2> 1
        custom:2> 4
        custom:2> 7
        custom:3> 2
        custom:3> 5
        custom:3> 8
         */
        // 自定义分区
        DataStreamSource<Integer> data = env.fromElements(1, 2, 3, 4, 5, 6, 7, 8);

        data.partitionCustom(new Partitioner<Integer>() {
            @Override
            public int partition(Integer key, int numPartitions) {
                return key % 3;
            }
        }, new KeySelector<Integer, Integer>() {
            @Override
            public Integer getKey(Integer value) throws Exception {
                return value;
            }
        }).print("custom").setParallelism(4);
        env.execute(TestPartition.class.getName());
    }
}
